import wx
import wx.wizard as wiz
import commands, os, re
import pexpect
import cPickle as pickle

from dmcal_utils.user_utils import root_check, get_user
from dmcal_utils.system_utils import get_free_loop_device, does_file_exist, size_converter, does_folder_exist

ID_NC_BROWSE_BUTTON = 1000

ID_EXIT = 100
ID_ABOUT = 200

class New_Container(wx.Dialog):
    def __init__(self, parent, id, title, add_container):
        wx.Dialog.__init__(self, parent, wx.ID_ANY, title, size=(400,275))
        
        self.add_container = add_container
        
        panel = wx.Panel(self, -1)
        vbox_master = wx.BoxSizer(wx.VERTICAL)
        
        wx.StaticText(panel, -1, "This will create a new container file.", (40,5))
        wx.StaticText(panel, -1, "Path: ", (5,35))
        self.pathTextBox = wx.TextCtrl(panel, -1, "", (50,30), size=(250,-1))
        self.pathTextBox.SetToolTip(wx.ToolTip("Please select or enter the\npath to the file you want to\nuse for your encrypted container"))
        self.pathbrowseButton = wx.Button(panel, ID_NC_BROWSE_BUTTON, "Browse", (300,30), size=(-1,30))
        wx.StaticText(panel, -1, "Size: ", (8, 65))
        self.sizeTextBox = wx.TextCtrl(panel, -1, "", (50,60), size=(250,-1))
        self.sizeTextBox.SetToolTip(wx.ToolTip("Please enter the size of the container in bytes."))
        
        wx.StaticText(panel, -1, "Method: ", (5, 95))
        
        self.urandomRadioButton = wx.RadioButton(panel, -1, "/dev/urandom", (90,93), style=wx.RB_GROUP)
        self.urandomRadioButton.SetToolTip(wx.ToolTip("/dev/urandom gives decent protection against\ncryptanalytic attacks and is only moderately slow."))
        self.randomRadioButton = wx.RadioButton(panel, -1, "/dev/random", (250,93))
        self.randomRadioButton.SetToolTip(wx.ToolTip("/dev/random gives much stronger protection against\ncryptanalytic attacks and is very slow."))
        
        wx.StaticText(panel, -1, "Password: ", (57,125))
        self.passwordTextBox = wx.TextCtrl(panel, -1, "", (143, 120), size=(250,-1), style=wx.TE_PASSWORD)
        wx.StaticText(panel, -1, "Re-Type: ", (73,155))
        self.repasswordTextBox = wx.TextCtrl(panel, -1, "", (143,150), size=(250, -1), style=wx.TE_PASSWORD)
        
        filesystems = ["ext3"]
        
        wx.StaticText(panel, -1, "Filesystem: ", (5,186))
        self.filesystemComboBox = (wx.ComboBox(panel, -1, pos=(104,180), size=(195,30), choices=filesystems, style=wx.CB_READONLY))
        self.filesystemComboBox.SetValue("ext3")
        
        freeloop = get_free_loop_device()
        wx.StaticText(panel, -1, "Loop Device: ", (5, 212))
        wx.StaticText(panel, -1, freeloop, (120,212))
        
        self.cancelButton = wx.Button(panel, wx.ID_CANCEL, "Cancel", (5,235))
        self.okButton = wx.Button(panel, wx.ID_OK, "OK", (300,235))

        self.progress = wx.Gauge(panel, range=100, pos=(100,240), size=(198,25))
        self.progress.Hide()

        self.Bind(wx.EVT_BUTTON, self.OnNCBrowse, id=ID_NC_BROWSE_BUTTON)
        self.Bind(wx.EVT_BUTTON, self.OnOK, id=wx.ID_OK)
        
    def OnOK(self, e):
        path = self.pathTextBox.GetValue()
        if not path:
            wx.MessageBox("The new containers path and filename must be given.", "Path Error", style=wx.OK | wx.ICON_ERROR)
            return 1
        else:
            filecheck = does_file_exist(path)
            pathcheck = does_folder_exist(path)
            if pathcheck:
                if filecheck:
                    wx.MessageBox("You can not create an encrypted\ncontainer in a file that already exists.", "File Error", style=wx.OK | wx.ICON_ERROR)
                    return 1
                else:
                    size = self.sizeTextBox.GetValue()
    
                    # Check for proper size entry
                    if re.match("^[1-9]([0-9]+)?(MB)$", size, re.I):                
                        size_in_bytes = size_converter(size, "by")
                        count = size_in_bytes / 1024
                    else:
                        wx.MessageBox("The size you entered was invalid.\n\nPlease enter the size using one of\nthe following:\n\nMB - Megabytes ", "Size Error", style=wx.OK | wx.ICON_ERROR)
                        return 0
              
                    if self.urandomRadioButton.GetValue() == True:
                        random_device = "/dev/urandom"
                    if self.randomRadioButton.GetValue() == True:
                        random_device = "/dev/random"
                
                    if password1:    
                        password1 = self.passwordTextBox.GetValue()
                        if password2:
                            password2 = self.repasswordTextBox.GetValue()
                        else:
                            wx.MessageBox("You must re-enter your password.", "Password Error", style=wx.OK | wx.ICON_ERROR)
                    else:
                        wx.MessageBox("You must enter a password.", "Password Error", style=wx.OK | wx.ICON_ERROR)
                        return 0
                
                if password1 == password2:
                    # Create the container file
                    dd_command = "dd if=%s of=%s bs=1024 count=%s" % (random_device, path, count)    
                    pexpect.run(dd_command)
            
                    # Bind the file to the free loop device
                    free_loop = get_free_loop_device()
                    bind_command = "losetup %s %s" % (free_loop, path)
                    pexpect.run(bind_command)
            
                    # Setup the encrypted filesystem
                    cryptsetup_command = """cryptsetup --verbose --verify-passphrase luksFormat %s""" % (free_loop)
                    child = pexpect.spawn(cryptsetup_command)
                    child.expect('.*')
                    child.send('YES\n')
                    child.expect('Enter LUKS passphrase: ')
                    child.sendline(password1)
                    child.expect('Verify passphrase: ')
                    child.sendline(password2)
                    result = child.readlines()
                                     
                    # Open the encrypted container
                    filename = os.path.basename(path)
                    open_command = "cryptsetup luksOpen %s %s" % (free_loop, filename)
                    child = pexpect.spawn(open_command)
                    child.expect("Enter LUKS passphrase: ")
                    child.sendline(password1)
                    result = child.readlines()
                 
                    # Format the internal filesystem
                    filesystem = self.filesystemComboBox.GetValue()
                    if filesystem == "ext3":
                        format_command = "mkfs.ext3 /dev/mapper/%s" % (filename)
                        pexpect.run(format_command)
                    
                    ### Cleanup ###
                 
                    # Close the container
                    close_command = "cryptsetup luksClose %s" % (filename)
                    pexpect.run(close_command)
                
                    # Delete the loop device
                    loop_delete_command = "losetup -d %s" % (free_loop)
                    pexpect.run(loop_delete_command)
                
                    self.add_container(filename, path)
                    self.Destroy()
            
                else:
                    wx.MessageBox("Passwords do not match.\nPlease try again.", "Password Error", style=wx.OK | wx.ICON_ERROR)
                    return 1
            else:            
                wx.MessageBox("The directory you are trying to use does not seem to exist.\nPlease correct this error and try again.", "Directory Error", style=wx.OK | wx.ICON_ERROR)
                return 1
                    
    def OnNCBrowse(self, e):
        dlg = wx.FileDialog(self, "Choose a file", os.getcwd(), "", "*.*", wx.OPEN)
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            mypath = os.path.basename(path)
            self.pathTextBox.WriteText(path)
        dlg.Destroy()
        
class MainWindow(wx.Frame):
    def __init__(self, parent, id, title):
        wx.Frame.__init__(self, parent, wx.ID_ANY, title, style=wx.MINIMIZE_BOX | wx.CLOSE_BOX)
        self.Bind(wx.EVT_CLOSE, self.OnExit)
        
        self.menubar = wx.MenuBar()
        
        self.container_locations = {}
        self.container_mounts = {}
        self.container_loops = {}
        
        self.ap = wx.ArtProvider()
        
        self.filemenu = wx.Menu()
        self.filemenu.AppendSeparator()
        
        self.quit_menu_item = wx.MenuItem(self.filemenu, ID_EXIT, "&Quit")
        self.quit_icon = self.ap.GetBitmap(wx.ART_QUIT,wx.ART_MENU, (16,16))
        self.quit_menu_item.SetBitmap(self.quit_icon)
        self.filemenu.AppendItem(self.quit_menu_item)
        
        
        self.addonsmenu = wx.Menu()
        
        self.helpmenu = wx.Menu()
        self.about_menu_item = wx.MenuItem(self.addonsmenu, ID_ABOUT, "&About")
        self.about_icon = self.ap.GetBitmap(wx.ART_QUESTION, wx.ART_MENU)
        self.about_menu_item.SetBitmap(self.about_icon)
        self.helpmenu.AppendItem(self.about_menu_item)
        
        self.menubar.Append(self.filemenu, "&File")
        self.menubar.Append(self.addonsmenu, "Add&ons")
        self.menubar.Append(self.helpmenu, "&Help")
        
        self.SetMenuBar(self.menubar)
        self.CreateStatusBar()
        
        # Button Layout
        self.sizer_main_buttons = wx.BoxSizer(wx.VERTICAL, )
        
        self.new_button = wx.Button(self, wx.ID_NEW)
        self.add_button = wx.Button(self, wx.ID_ADD)
        self.remove_button = wx.Button(self, wx.ID_REMOVE)
        self.mount_button = wx.Button(self, wx.ID_OPEN)
        self.unmount_button = wx.Button(self, wx.ID_CLOSE)
        self.exit_button = wx.Button(self, wx.ID_EXIT)
        
        self.sizer_main_buttons.Add(self.new_button, 0, wx.EXPAND)
        self.sizer_main_buttons.Add(self.add_button, 0, wx.EXPAND)
        self.sizer_main_buttons.Add(self.remove_button, 0, wx.EXPAND)
        self.sizer_main_buttons.Add(self.mount_button, 0, wx.EXPAND)
        self.sizer_main_buttons.Add(self.unmount_button, 0, wx.EXPAND)
        self.sizer_main_buttons.Add(self.exit_button, 0, wx.EXPAND)
        
        # Main Layout
        self.sizer_main = wx.BoxSizer(wx.HORIZONTAL)
        
        # Setup the list
        self.container_list = wx.ListCtrl(self, -1, size = (600,100), style = wx.LC_REPORT)
        self.container_list.InsertColumn(0, "", width = 25)
        self.container_list.InsertColumn(1, "Container", width = 150)
        self.container_list.InsertColumn(2, "Mount", width = 425)
                
        self.image_list = wx.ImageList(16,16)
        self.container_list.SetImageList(self.image_list, wx.IMAGE_LIST_SMALL)
        
        self.sizer_main.Add(self.container_list, 0, wx.EXPAND)
        self.sizer_main.Add(self.sizer_main_buttons, 0, wx.EXPAND)         
        
        # Menu Events
        wx.EVT_MENU(self, ID_EXIT, self.OnExit)
        wx.EVT_MENU(self, ID_ABOUT, self.OnAbout)
        
        # Button Events
        self.Bind(wx.EVT_BUTTON, self.OnExit, id=wx.ID_EXIT)
        self.Bind(wx.EVT_BUTTON, self.OnNew, id=wx.ID_NEW)
        self.Bind(wx.EVT_BUTTON, self.OnAdd, id=wx.ID_ADD)
        self.Bind(wx.EVT_BUTTON, self.OnRemove, id=wx.ID_REMOVE)
        self.Bind(wx.EVT_BUTTON, self.OnMount, id=wx.ID_OPEN)
        self.Bind(wx.EVT_BUTTON, self.OnUnMount, id=wx.ID_CLOSE)
        
        # List Control Events
        self.Bind(wx.EVT_LIST_ITEM_SELECTED, self.OnContainerSelected, self.container_list)
        
        # Check to see if the program was run as root
        root_status = root_check()
        if root_status == "no":
            self.menubar.Enable(ID_NEW_CONTAINER, False)
            self.new_button.Enable(False)
            self.mount_button.Enable(False)
            self.unmount_button.Enable(False)
            self.SetStatusText("Program running in limited mode.  Root permissions and path needed.")
                
        self.SetSizer(self.sizer_main)
        self.SetAutoLayout(1)
        self.sizer_main.Fit(self)
        self.CenterOnScreen()
        self.Show(1)
        
        self.OnLoad()
        
    def OnUnMount(self, e):
        selected_count = self.container_list.GetSelectedItemCount()
        if selected_count > 0:
            selected_item = self.container_list.GetFocusedItem()
            item = self.container_list.GetItem(selected_item, 2)
            mount_path = item.GetText()
            if mount_path == "":
                wx.MessageBox("This container does not have a mount point.", "File Error", style=wx.OK | wx.ICON_ERROR)
            else:
                selected_item = self.container_list.GetFocusedItem()
                item = self.container_list.GetItem(selected_item, 1)
                filename = item.GetText()
                
                # Unmount the container
                unmount_command = "umount %s" % (mount_path)
                pexpect.run(unmount_command)
                
                # Close the encrypted container
                close_command = "cryptsetup luksClose %s" % (filename)
                pexpect.run(close_command)
                
                # Delete the free'd up loop device
                free_loop = self.container_loops[filename]
                delete_command = "losetup -d %s" % (free_loop)
                pexpect.run(delete_command)
                
                del self.container_loops[filename]
                
                # Update icon
                selected_item = self.container_list.GetFocusedItem()
                self.mark = self.image_list.Add(self.ap.GetBitmap(wx.ART_CROSS_MARK,wx.ART_FRAME_ICON, (16,16)))        
                self.container_list.SetItemImage(selected_item, self.mark)
                
                # Clear mount text
                self.container_list.SetStringItem(selected_item,2,"")
        else:
            wx.MessageBox("You must select a mounted container.", "File Error", style=wx.OK | wx.ICON_ERROR)
        
    def OnMount(self, e):
        selected_count  = self.container_list.GetSelectedItemCount()
        if selected_count > 0:
            selected_item = self.container_list.GetFocusedItem()
            item = self.container_list.GetItem(selected_item, 2)
            mount_path = item.GetText()
            if mount_path == "":
                dlg = wx.DirDialog(self, "Choose a directory:", style=wx.DD_DEFAULT_STYLE | wx.DD_NEW_DIR_BUTTON)
                if dlg.ShowModal() == wx.ID_OK:
                    mount_path = dlg.GetPath()
                    self.container_list.SetStringItem(selected_item,2,mount_path)
                    item = self.container_list.GetItem(selected_item, 1)
                    filename = item.GetText()
                    self.container_mounts[filename] = mount_path
                dlg.Destroy()
                
                self.Mount_Container(filename)
            else:
                item = self.container_list.GetItem(selected_item, 1)
                filename = item.GetText()
            
                self.Mount_Container(filename)
        else:
            wx.MessageBox("You must have a container file selected.\n\nNOTE: If there is only on container in\nthe list it will be automatically selected.", "File Error", style=wx.OK | wx.ICON_ERROR)
               
    def Mount_Container(self, filename):
        free_loop = get_free_loop_device()
        container_path = self.container_locations[filename]
        mount_path = self.container_mounts[filename]
        
        # Attach the container to the loop device
        attach_command = "losetup %s %s" % (free_loop,container_path)
        pexpect.run(attach_command)
        
        # Prompt for password
        dlg = wx.TextEntryDialog(self, 'Please enter your password for the container.','Enter Your Password')
        dlg.SetValue("")
        if dlg.ShowModal() == wx.ID_OK:
            password = dlg.GetValue()
        else: return 1
        dlg.Destroy()
        
        # Open the encrypted container
        open_command = "cryptsetup luksOpen %s %s" % (free_loop, filename)
        child = pexpect.spawn(open_command, timeout=10)
        child.expect("Enter LUKS passphrase: ")
        child.sendline(password)
        try:
            result = child.readlines()
        except pexpect.TIMEOUT:
            wx.MessageBox("There was a problem opening the container.\nPlease check your password and try again.", "Error Opening Container", style=wx.OK | wx.ICON_ERROR)
            
            ### Cleanup ###
            
            # Un-hook the container from the loop device.
            unhook_command = "losetup -d %s" % (free_loop)
            pexpect.run(unhook_command)
                        
            return 1
        
        # Mount the container to the folder
        mount_command = "mount /dev/mapper/%s %s" % (filename, mount_path)
        pexpect.run(mount_command)
        
        self.container_loops[filename] = free_loop
        
        # Update icon
        selected_item = self.container_list.GetFocusedItem()
        self.tick = self.image_list.Add(self.ap.GetBitmap(wx.ART_TICK_MARK,wx.ART_FRAME_ICON, (16,16)))        
        self.container_list.SetItemImage(selected_item, self.tick)   
                          
    def OnRemove(self, e):
        selected_item = self.container_list.GetFocusedItem()
        item = self.container_list.GetItem(selected_item,1)
        filename = item.GetText()
        path = self.container_locations[filename]
        text = "Removed (%s) %s" % (filename,path)
        self.container_list.DeleteItem(selected_item)
        try:
            del self.container_locations[filename]
            del self.container_mounts[filename]
        except:
            pass
        self.SetStatusText(text)
       
    def OnAdd(self, e):
        dlg = wx.FileDialog(self, "Choose a file", os.getcwd(), "", "*.*", wx.OPEN)
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            filename = os.path.basename(path)
            self.add_container(filename, path)
        dlg.Destroy()
        
    def OnNew(self, e):
         new_container = New_Container(None, -1, "New Container", self.add_container)
         new_container.ShowModal()
                    
    def OnExit(self, e):
        really_quit = wx.MessageDialog(None, 'Are you sure to quit?', 'Really Quit?', 
            wx.YES_NO | wx.NO_DEFAULT | wx.ICON_QUESTION)
        
        if really_quit.ShowModal() == wx.ID_YES:        
            save_file = self.dmcal_save+"/.save"
            
            list_count = self.container_list.GetItemCount()
            data_store = {}
            i = 0
            
            if does_file_exist(save_file):
                os.remove(save_file)
                os.mknod(save_file, 0700)

                while i < list_count:
                    temp_store = {}
                    #raw_icon = self.container_list.GetItem(i, 0)
                    #icon = raw_icon.GetImage()
                    raw_name = self.container_list.GetItem(i, 1)
                    name = raw_name.GetText()
                    location = self.container_locations[name]
                    raw_mount = self.container_list.GetItem(i, 2)
                    mount = raw_mount.GetText()
                    #temp_store["icon"] = icon
                    temp_store["name"] = name
                    temp_store["location"] = location
                    try:
                        loop = self.container_loops[name]
                        temp_store["loop"] = loop
                    except:
                        temp_store["loop"] = ""
                    temp_store["mount"] = mount
                    data_store[name] = temp_store                    
                    i = i + 1
                
                saved_file = open(save_file, "w")
                pickle.dump(data_store, saved_file)
                saved_file.close()                   
            else:
                os.mknod(save_file, 0700)
                
                while i < list_count:
                    temp_store = {}
                    #raw_icon = self.container_list.GetItem(i, 0)
                    #icon = raw_icon.GetImage()
                    raw_name = self.container_list.GetItem(i, 1)
                    name = raw_name.GetText()
                    location = self.container_locations[name]
                    raw_mount = self.container_list.GetItem(i, 2)
                    mount = raw_mount.GetText()
                    #temp_store["icon"] = icon
                    temp_store["name"] = name
                    temp_store["location"] = location
                    try:
                        loop = self.container_loops[name]
                        temp_store["loop"] = loop
                    except:
                        temp_store["loop"] = ""
                    temp_store["mount"] = mount
                    data_store[name] = temp_store                    
                    i = i + 1
                
                saved_file = open(save_file, 'w')
                pickle.dump(data_store, saved_file)
                saved_file.close()
                
            self.Destroy()
        
    def OnAbout(self, e):
        
        description = """DMCaL is, at its most basic level, a program to help users create and maintain encrypted container files.  It also will handle mounting and un-mounting along with any maintenance that needs to be done such as changes in keys."""
        license = """This is the license information area."""
        
        info = wx.AboutDialogInfo()
        # info.SetIcon(wx.Icon('icons/hunter.png', wx.BITMAP_TYPE_PNG))
        info.SetName('DMCaL')
        info.SetVersion('alpha')
        info.SetDescription(description)
        info.SetCopyright('(c) 2008 Nathan McBride')
        info.SetWebSite('http://www.nombyte.com')
        info.SetLicense(license)
        info.AddDeveloper('Nathan McBride')
        info.AddDocWriter('Nathan McBride')
        info.AddArtist('Nathan McBride')
        info.AddTranslator('Nathan McBride')
        
        wx.AboutBox(info)
    
    def add_container(self, name, path):
        list_count = self.container_list.GetItemCount()

        self.mark = self.image_list.Add(self.ap.GetBitmap(wx.ART_CROSS_MARK,wx.ART_FRAME_ICON, (16,16)))        
        self.container_list.InsertImageItem(list_count, self.mark)
                
        self.container_list.SetStringItem(list_count, 1, name)
        self.container_locations[name]=path
        text = "Added: (%s) %s" % (name,path)
        self.SetStatusText(text)
        
    def add_container_with_mount(self, name, path, mount, loop_dev):
        list_count = self.container_list.GetItemCount()
    
        self.mark = self.image_list.Add(self.ap.GetBitmap(wx.ART_TICK_MARK,wx.ART_FRAME_ICON, (16,16)))        
        self.container_list.InsertImageItem(list_count, self.mark)
        self.container_list.SetStringItem(list_count, 1, name)
        self.container_list.SetStringItem(list_count, 2, mount)
        self.container_locations[name] = path
        self.container_mounts[name] = mount
        self.container_loops[name] = loop_dev
        
        text = "Added: (%s) %s" % (name,path)
        self.SetStatusText(text)
    
    def OnContainerSelected(self, e):
        pass
    
    def OnLoad(self):
        # Check if save folder exists
        self.user = get_user()
        self.user_home =  os.path.expanduser('~')
        self.dmcal_save = self.user_home + "/.dmcal"
        if not does_folder_exist(self.dmcal_save):
            os.mkdir(self.dmcal_save, 0700)
        else:
            # Check if save file exists
            saved_file = self.dmcal_save+"/.save"
            if does_file_exist(saved_file):
                # Load the data
                load_file = open(saved_file, 'r')
                try:
                    loaded_data_store = pickle.load(load_file)
                except:
                    return 1
                # Add the data to the list
                for container in loaded_data_store:
                    temp_data = loaded_data_store[container]
                    if temp_data["mount"] == "":
                        self.add_container(temp_data["name"], temp_data["location"])
                    else:
                        self.add_container_with_mount(temp_data["name"], temp_data["location"], temp_data["mount"], temp_data["loop"])
                    
                

app = wx.App()
frame = MainWindow(None, -1, "DMCaL")
app.MainLoop()
